package com.ag.improvedjava.iterable;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.TreeMap;

import com.ag.improvedjava.utils.Grouping;
import com.ag.improvedjava.utils.Lookup;

final class LookupImpl<K, V> implements Lookup<K, V> {

	private final Map<K, List<V>> map;
	private final List<K> keys;

	LookupImpl() {
		map = new TreeMap<K, List<V>>();
		keys = new ArrayList<K>();
	}
	
	LookupImpl(Comparator<K> comparer) {
		map = new TreeMap<K, List<V>>(comparer);
		keys = new ArrayList<K>();
	}

	void add(K key, V value) {
		List<V> list = map.get(key);
		if (list == null) {
			list = new ArrayList<V>();
			map.put(key, list);
			keys.add(key);
		}
		list.add(value);
	}

	@Override
	public int count() {
		return map.size();
	}

	@Override
	public Iterable<V> values(K key) {
		return contains(key) ? new ArrayList<>(map.get(key)) : Collections
				.<V> emptyList();
	}

	@Override
	public boolean contains(K key) {
		return map.containsKey(key);
	}

	@Override
	public Iterator<Grouping<K, V>> iterator() {
		return new Iterator<Grouping<K, V>>() {
			private int count = 0;

			@Override
			public boolean hasNext() {
				return count < keys.size();
			}

			@Override
			public Grouping<K, V> next() {
				if (!hasNext()) {
					throw new NoSuchElementException();
				}
				K key = keys.get(count++);
				return new GroupingImpl<>(key, map.get(key));
			}

			@Override
			public void remove() {
				throw new UnsupportedOperationException();
			}
		};
	}

}
